"""
@Desc: 均线交叉策略的流程
数据准备：获取股票或金融产品的历史价格数据。
计算均线：计算短期均线和长期均线。
生成信号：根据均线交叉生成买卖信号。
回测：模拟交易，计算策略的收益和表现。
可视化：绘制价格、均线和买卖信号。

@Auth: meihongliang-m2
@Date: 2025/3/18-15:56
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf  # 用于获取股票数据


# 1. 获取数据
def get_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data['Adj Close']  # 使用调整后的收盘价


# 2. 计算均线
def calculate_moving_averages(data, short_window, long_window):
    data['Short_MA'] = data['Adj Close'].rolling(window=short_window, min_periods=1).mean()
    data['Long_MA'] = data['Adj Close'].rolling(window=long_window, min_periods=1).mean()
    return data


# 3. 生成交易信号
def generate_signals(data):
    data['Signal'] = 0  # 初始化信号列
    data['Signal'][short_window:] = np.where(data['Short_MA'][short_window:] > data['Long_MA'][short_window:], 1, -1)
    data['Position'] = data['Signal'].diff()  # 计算仓位变化（买入或卖出）
    return data


# 4. 回测策略
def backtest(data):
    initial_capital = 10000  # 初始资金
    position = 0  # 持仓数量
    portfolio_value = []  # 记录每日投资组合价值

    for i in range(len(data)):
        if data['Position'][i] == 1:  # 买入信号
            position = initial_capital // data['Adj Close'][i]  # 计算买入数量
            initial_capital -= position * data['Adj Close'][i]  # 更新剩余资金
        elif data['Position'][i] == -1:  # 卖出信号
            initial_capital += position * data['Adj Close'][i]  # 更新资金
            position = 0  # 清空持仓

        portfolio_value.append(initial_capital + position * data['Adj Close'][i])  # 计算每日组合价值

    data['Portfolio_Value'] = portfolio_value
    return data


# 5. 可视化结果
def plot_results(data):
    plt.figure(figsize=(14, 7))
    plt.plot(data['Adj Close'], label='Price', alpha=0.5)
    plt.plot(data['Short_MA'], label='Short MA (5)', alpha=0.75)
    plt.plot(data['Long_MA'], label='Long MA (20)', alpha=0.75)
    plt.scatter(data.index, data['Adj Close'], c=data['Position'], cmap='viridis', label='Buy/Sell Signal', marker='^',
                s=100)
    plt.title('Moving Average Crossover Strategy')
    plt.legend()
    plt.show()


# 主程序
if __name__ == "__main__":
    # 参数设置
    ticker = 'AAPL'  # 股票代码
    start_date = '2020-01-01'
    end_date = '2023-01-01'
    short_window = 5  # 短期均线窗口
    long_window = 20  # 长期均线窗口

    # 获取数据
    data = get_data(ticker, start_date, end_date)
    data = pd.DataFrame(data)
    # 计算均线
    data = calculate_moving_averages(data, short_window, long_window)
    # 生成信号
    data = generate_signals(data)
    # 回测策略
    data = backtest(data)
    # 可视化结果
    plot_results(data)
    # 打印最终收益
    final_portfolio_value = data['Portfolio_Value'].iloc[-1]
    print(f"Final Portfolio Value: {final_portfolio_value:.2f}")
